# vectordb.py
import chromadb

def get_collection(db_path: str, collection_name: str = "default"):
    chromadb_client = chromadb.PersistentClient(db_path)
    cdb_collection = chromadb_client.get_or_create_collection(name=collection_name)
    return cdb_collection

def store_to_db(db_path: str, data_embedding: list[list[float]], chunks: list[str]):
    cdb_collection = get_collection(db_path)
    
    ids = [str(i) for i in range(len(data_embedding))]
    cdb_collection.add(ids=ids, embeddings=data_embedding, documents=chunks)
    
def query_from_db(db_path:str,query_embedding:list[list[float]],top_k:int):
    cdb_collection = get_collection(db_path)
    return cdb_collection.query(query_embeddings=query_embedding, n_results=top_k)